# include "SaveFileFilter.h"

extern "C" {
# include "libavformat/avformat.h"
# include "libavcodec/avcodec.h"
}

SaveFileFilter::SaveFileFilter()
{
	
}

bool SaveFileFilter::initFileFilter(const char* _file_name, uint8_t* _video_encoder_codec_ctx, uint8_t* _audio_encoder_codec_ctx)
{
	if (avformat_alloc_output_context2(&this->m_save_ctx, nullptr, nullptr, _file_name) < 0) {
		// TODO log
		return false;
	}
	if (avio_open(&this->m_save_ctx->pb, _file_name, AVIO_FLAG_WRITE | AVIO_FLAG_READ) < 0) {
		// TODO log

		avformat_free_context(this->m_save_ctx);
		this->m_save_ctx = nullptr;
		return false;
	}
	
	if (_video_encoder_codec_ctx != nullptr) {
		AVCodecContext* video_encoder_codec_ctx = reinterpret_cast<AVCodecContext*>(_video_encoder_codec_ctx);
		this->m_save_video_stream = avformat_new_stream(this->m_save_ctx, nullptr);
		avcodec_parameters_from_context(this->m_save_video_stream->codecpar, video_encoder_codec_ctx);
		this->m_video_time_base = new AVRational;
		this->m_video_time_base->den = video_encoder_codec_ctx->time_base.den;
		this->m_video_time_base->num = video_encoder_codec_ctx->time_base.num;
		this->m_video_pts = 0;
	}

	if (_audio_encoder_codec_ctx != nullptr) {
		AVCodecContext* audio_encoder_codec_ctx = reinterpret_cast<AVCodecContext*>(_audio_encoder_codec_ctx);
		this->m_save_audio_stream = avformat_new_stream(this->m_save_ctx, nullptr);
		avcodec_parameters_from_context(this->m_save_audio_stream->codecpar, audio_encoder_codec_ctx);
		this->m_audio_pts = 0;
	}

	// av_dump_format(this->m_saveCtx, 0, _save_file_name, 1); TODO dump to log
	avformat_write_header(this->m_save_ctx, nullptr);
	this->m_thread = new std::thread(&SaveFileFilter::writeFrame, this);
	return true;
}

void SaveFileFilter::saveFile()
{
	this->m_ready = false;
	if (this->m_save_ctx == nullptr) {
		return;
	}
	this->m_cv.notify_one();
	this->m_thread->join();
	av_write_trailer(this->m_save_ctx);
	avio_close(this->m_save_ctx->pb);
	avformat_free_context(this->m_save_ctx);
	this->m_save_ctx = nullptr;
	AVPacket* _packet = nullptr;
	while (!this->m_queue.empty()) {
		_packet = m_queue.front();
		av_packet_free(&_packet);
		m_queue.pop();
	}
}

void SaveFileFilter::afterVideoEncode(AVPacket* _packet)
{
	if (this->m_ready) {
		_packet->stream_index = this->m_save_video_stream->index;
		_packet->dts = _packet->pts = this->m_video_pts;
		av_packet_rescale_ts(_packet, *this->m_video_time_base, this->m_save_video_stream->time_base);
		this->m_video_pts += 1;
		this->push(_packet);
	}
	EncoderFilter::afterVideoEncode(_packet);
}

void SaveFileFilter::afterAudioEncode(AVPacket* _packet)
{
	if (this->m_ready) {
		_packet->stream_index = this->m_save_audio_stream->index;
		_packet->dts = _packet->pts = this->m_audio_pts;
		av_packet_rescale_ts(_packet, this->m_save_audio_stream->time_base, this->m_save_audio_stream->time_base);
		this->m_audio_pts += 1024;
		this->push(_packet);
	}
	EncoderFilter::afterAudioEncode(_packet);
}

SaveFileFilter::~SaveFileFilter()
{
	this->saveFile();
	delete this->m_video_time_base;
}

void SaveFileFilter::writeFrame()
{
	this->m_ready = true;
	while (this->m_ready) {
		AVPacket* _packet = nullptr;
		{
			std::unique_lock<std::mutex> lock(this->m_mutex);
			this->m_cv.wait(lock, [this]() {
				return !this->m_queue.empty() || !this->m_ready;
			});
			if (!this->m_ready) {
				break;
			}
			_packet = this->m_queue.front();
			this->m_queue.pop();
		}
		if (_packet != nullptr) {
			av_interleaved_write_frame(this->m_save_ctx, _packet);
			av_packet_free(&_packet);
		}
	}
	
}

void SaveFileFilter::push(AVPacket* _packet)
{
	AVPacket* dst = av_packet_alloc();
	av_new_packet(dst, _packet->size);
	memcpy_s(dst->data, _packet->size, _packet->data, _packet->size);
	av_packet_copy_props(dst, _packet);
	std::lock_guard<std::mutex> lock(this->m_mutex);
	this->m_queue.push(dst);
	this->m_cv.notify_one();
}
